Models Documentation
This documentation provides an overview of key models, focusing on essential input parameters and high-level functionality.
Conditional Entropy Models
- class encoding_information.models.conditional_entropy_models.AnalyticGaussianNoiseModel(sigma)
Bases:
MeasurementNoiseModelAnalytical model for estimating the conditional entropy H(Y | X) when the noise proces is additive independent Gaussian noise at each pixel.
Initialize the AnalyticGaussianNoiseModel.
- Parameters:
sigma (float) – Standard deviation of the Gaussian noise.
- __init__(sigma)
Initialize the AnalyticGaussianNoiseModel.
- Parameters:
sigma (float) – Standard deviation of the Gaussian noise.
- estimate_conditional_entropy(images=None)
Compute the conditional entropy H(Y | X) for Gaussian noise.
- Parameters:
images (jax.Array, optional) – Unused in this model. The parameter is kept for compatibility.
- Returns:
The computed conditional entropy, given by the formula: 0.5 * log(2 * pi * e * sigma^2).
- Return type:
float
- class encoding_information.models.conditional_entropy_models.PoissonNoiseModel
Bases:
MeasurementNoiseModelPoisson noise model for estimating the conditional entropy H(Y | X) from empirical data.
- estimate_conditional_entropy(images)
Compute the conditional entropy H(Y | X) for Poisson noise.
- Parameters:
images (jax.Array) – A dataset of images, preferably clean. This is used to compute the conditional entropy.
- Returns:
The average conditional entropy per pixel, computed using a Gaussian approximation.
- Return type:
float
Gaussian Process Models
- class encoding_information.models.gaussian_process.StationaryGaussianProcess(images, eigenvalue_floor=0.001, seed=None, verbose=False)
Bases:
MeasurementModelStationary 2D Gaussian process for single-channel images.
This model learns a stationary covariance matrix and uses it to generate or evaluate samples.
Initialize the StationaryGaussianProcess model and estimate the covariance matrix.
- Parameters:
images (ndarray) – Input images used to estimate the stationary covariance matrix.
eigenvalue_floor (float, optional) – The minimum eigenvalue to ensure the covariance matrix is positive definite (default is 1e-3).
seed (int, optional) – Random seed for reproducibility.
verbose (bool, optional) – Whether to print progress and additional information (default is False).
- __init__(images, eigenvalue_floor=0.001, seed=None, verbose=False)
Initialize the StationaryGaussianProcess model and estimate the covariance matrix.
- Parameters:
images (ndarray) – Input images used to estimate the stationary covariance matrix.
eigenvalue_floor (float, optional) – The minimum eigenvalue to ensure the covariance matrix is positive definite (default is 1e-3).
seed (int, optional) – Random seed for reproducibility.
verbose (bool, optional) – Whether to print progress and additional information (default is False).
- compute_negative_log_likelihood(images, data_seed=None, verbose=True, seed=None, average=True)
Compute the negative log-likelihood of the provided images under the learned Gaussian process.
- Parameters:
images (ndarray) – Input images to evaluate.
data_seed (int, optional) – Random seed for shuffling the data.
verbose (bool, optional) – Whether to print progress (default is True).
seed (int, optional, deprecated) – Deprecated argument for random seed, use data_seed instead.
average (bool) – If true, return the average NLL, if false return per-example
- Returns:
The negative log-likelihood of the provided images.
- Return type:
float
- fit(train_images=None, data_seed=None, learning_rate=100.0, max_epochs=60, steps_per_epoch=1, patience=15, batch_size=12, num_val_samples=None, percent_samples_for_validation=0.1, eigenvalue_floor=0.001, gradient_clip=1, momentum=0.9, precondition_gradient=False, verbose=True)
Train the Stationary Gaussian Process model on the provided images.
- Parameters:
train_images (ndarray, optional) – The training dataset consisting of images (default is the images provided during initialization).
data_seed (int, optional) – Random seed for shuffling the data.
learning_rate (float, optional) – Learning rate for optimization (default is 1e2).
max_epochs (int, optional) – Maximum number of training epochs (default is 60).
steps_per_epoch (int, optional) – Number of steps per epoch (default is 1).
patience (int, optional) – Number of epochs to wait for early stopping (default is 15).
batch_size (int, optional) – The number of images in each batch (default is 12).
num_val_samples (int, optional) – Number of validation samples (default is computed automatically based on percent_samples_for_validation).
percent_samples_for_validation (float, optional) – Fraction of samples to use for validation (default is 0.1).
eigenvalue_floor (float, optional) – The minimum eigenvalue to ensure the covariance matrix is positive definite (default is 1e-3).
gradient_clip (float, optional) – Maximum value to clip the gradient to avoid large updates (default is 1).
momentum (float, optional) – Momentum parameter for the optimizer (default is 0.9).
precondition_gradient (bool, optional) – Whether to precondition the gradient using the Fisher Information Matrix (default is False).
verbose (bool, optional) – Whether to print progress (default is True).
- Returns:
val_loss_history – A list of validation losses for each epoch.
- Return type:
list
- generate_samples(num_samples, sample_shape=None, ensure_nonnegative=True, seed=None, verbose=True)
Generate new image samples from the learned Gaussian process.
- Parameters:
num_samples (int) – The number of samples to generate.
sample_shape (tuple, optional) – Shape of the samples to generate (default is the shape of the training images).
ensure_nonnegative (bool, optional) – Whether to ensure all pixel values are non-negative (default is True).
seed (int, optional) – Random seed for reproducibility.
verbose (bool, optional) – Whether to print progress (default is True).
- Returns:
Generated image samples.
- Return type:
ndarray
- class encoding_information.models.gaussian_process.FullGaussianProcess(data, eigenvalue_floor=0.001, seed=None, verbose=False, add_uniform_noise=True)
Bases:
MeasurementModelFull (non-stationary) Gaussian process for arbitrary data.
This model estimates a full covariance matrix from the data and uses it to generate or evaluate samples.
Initialize the Full Gaussian Process model and estimate the mean and covariance matrix from the data.
- Parameters:
data (ndarray) – Input data used to estimate the Gaussian process.
eigenvalue_floor (float, optional) – Minimum eigenvalue to ensure the covariance matrix is positive definite (default is 1e-3).
seed (int, optional) – Random seed for reproducibility.
verbose (bool, optional) – Whether to print progress during computation (default is False).
add_uniform_noise (bool, optional) – For discrete-valued data, add uniform noise (default True). For continuous-valued data, set to False.
- __init__(data, eigenvalue_floor=0.001, seed=None, verbose=False, add_uniform_noise=True)
Initialize the Full Gaussian Process model and estimate the mean and covariance matrix from the data.
- Parameters:
data (ndarray) – Input data used to estimate the Gaussian process.
eigenvalue_floor (float, optional) – Minimum eigenvalue to ensure the covariance matrix is positive definite (default is 1e-3).
seed (int, optional) – Random seed for reproducibility.
verbose (bool, optional) – Whether to print progress during computation (default is False).
add_uniform_noise (bool, optional) – For discrete-valued data, add uniform noise (default True). For continuous-valued data, set to False.
- compute_negative_log_likelihood(data, data_seed=None, verbose=True, seed=None, average=True)
Compute the negative log-likelihood of the provided data under the Gaussian process.
- Parameters:
data (ndarray) – Input data to evaluate.
data_seed (int, optional) – Random seed for shuffling the data.
verbose (bool, optional) – Whether to print progress (default is True).
seed (int, optional, deprecated) – Deprecated argument for random seed, use data_seed instead.
average (bool, optional) – Whether to average the negative log-likelihood over all samples (default is True).
- Returns:
The negative log-likelihood of the provided data.
- Return type:
float
- generate_samples(num_samples, sample_shape=None, ensure_nonnegative=True, seed=None, verbose=True)
Generate new samples from the learned Gaussian process.
- Parameters:
num_samples (int) – Number of samples to generate.
sample_shape (tuple, optional) – Shape of the samples to generate (default is the shape of the training data).
ensure_nonnegative (bool, optional) – Whether to ensure all values are non-negative (default is True).
seed (int, optional) – Random seed for reproducibility.
verbose (bool, optional) – Whether to print progress (default is True).
- Returns:
Generated samples.
- Return type:
ndarray
Pixel CNN Model
- class encoding_information.models.pixel_cnn.PixelCNN(num_hidden_channels=64, num_mixture_components=40)
Bases:
MeasurementModelThe PixelCNN model for autoregressive image modeling.
This class handles the training and evaluation of the PixelCNN model and wraps the Flax implementation in a higher-level interface that conforms to the MeasurementModel class. It provides methods for fitting the model to data, computing the negative log-likelihood of images, and generating new images.
The number of hidden channels in the model.
- Type:
int
- num_mixture_components
The number of components in the mixture density output.
- Type:
int
Initialize the PixelCNN model with image shape, number of hidden channels, and mixture components.
- Parameters:
num_hidden_channels (int) – Number of hidden channels in the convolutional layers.
num_mixture_components (int) – Number of mixture components for the output layer.
- __init__(num_hidden_channels=64, num_mixture_components=40)
Initialize the PixelCNN model with image shape, number of hidden channels, and mixture components.
- Parameters:
num_hidden_channels (int) – Number of hidden channels in the convolutional layers.
num_mixture_components (int) – Number of mixture components for the output layer.
- compute_negative_log_likelihood(data, conditioning_vecs=None, data_seed=None, average=True, verbose=True, seed=None)
Compute the negative log-likelihood (NLL) of images under the trained PixelCNN model.
- Parameters:
data (ndarray) – The input images for which to compute the NLL.
conditioning_vecs (ndarray, optional) – Vectors to condition the image generation process (e.g., class labels).
data_seed (int, optional) – Seed for data shuffling.
average (bool, optional) – If True, return the average NLL over all images (default is True).
verbose (bool, optional) – Whether to print progress (default is True).
seed (int, optional) – Deprecated. Use data_seed instead.
- Returns:
nll – The negative log-likelihood of the input images.
- Return type:
float
- fit(train_images, condition_vectors=None, learning_rate=0.01, max_epochs=200, steps_per_epoch=100, patience=40, sigma_min=1, batch_size=64, num_val_samples=None, percent_samples_for_validation=0.1, do_lr_decay=False, verbose=True, add_gaussian_noise=False, add_uniform_noise=True, model_seed=None, data_seed=None, use_positional_embedding=False, seed=None)
Train the PixelCNN model on a dataset of images.
- Parameters:
train_images (ndarray) – The input dataset, with shape (N, H, W, C).
condition_vectors (ndarray, optional) – Vectors to condition the image generation process (e.g., class labels).
learning_rate (float, optional) – The learning rate for optimization (default is 1e-2).
max_epochs (int, optional) – The maximum number of training epochs (default is 200).
steps_per_epoch (int, optional) – The number of steps per epoch (default is 100).
patience (int, optional) – The number of epochs to wait before early stopping (default is 40).
sigma_min (float, optional) – The minimum standard deviation for the mixture density output (default is 1).
batch_size (int, optional) – The batch size for training (default is 64).
num_val_samples (int, optional) – The number of validation samples. If None, a percentage is used (default is None).
percent_samples_for_validation (float, optional) – The percentage of samples to use for validation (default is 0.1).
do_lr_decay (bool, optional) – Whether to apply learning rate decay during training (default is False).
verbose (bool, optional) – Whether to print progress during training (default is True).
add_gaussian_noise (bool, optional) – Whether to add Gaussian noise to the training images (default is False).
add_uniform_noise (bool, optional) – Whether to add uniform noise to the training images (default is True).
model_seed (int, optional) – Seed for model initialization.
data_seed (int, optional) – Seed for data shuffling.
- Returns:
val_loss_history – A list of validation loss values for each epoch.
- Return type:
list
- generate_samples(num_samples, conditioning_vecs=None, sample_shape=None, ensure_nonnegative=True, seed=None, verbose=True)
Generate new images from the trained PixelCNN model by sampling pixel by pixel.
- Parameters:
num_samples (int) – Number of images to generate.
conditioning_vecs (jax.Array, optional) – Optional conditioning vectors. If provided, the shape should match (num_samples, condition_vector_size). Default is None.
sample_shape (tuple of int or int, optional) – Shape of the images to generate. If None, the model’s image_shape is used. If a single int is provided, it will be treated as a square shape. Default is None.
ensure_nonnegative (bool, optional) – If True, ensure that the generated pixel values are non-negative. Default is True.
seed (int, optional) – Random seed for reproducibility. Default is 123 if not provided.
verbose (bool, optional) – If True, display progress during the generation process. Default is True.
- Returns:
Generated images with the specified shape.
- Return type:
jax.Array